
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <CL/cl.h>
#include <jni.h> 

const char *KernelSource = "\n" \
"__kernel void square(                                                   \n" \
"   __global float* val,                                                 \n" \
"   __global float* y,                                                   \n" \
"	__global int* row,	 												 \n" \
"   __global int* col,                                              	 \n" \
"   __global float* x,                                            		 \n" \
"   const unsigned int NUM_ITERATIONS,                                   \n" \
"   const unsigned int nthreads,                                       	 \n" \
"   __global int* lowsum,                                             	 \n" \
"   __global int* highsum,                                             	 \n" \
"   __local int* loclsum,                                             	 \n" \
"   __local int* lochsum)                                             	 \n" \
"{                                                                     	 \n" \
"int id = get_global_id(0);												 \n" \
"int lid = get_local_id(0);												 \n" \
"int lsize = get_local_size(0);												 \n" \
"int gid = get_group_id(0);												 \n" \
"if(id<nthreads){ 							 \n" \
"	int i;					 											 \n" \
"	int j;					 											 \n" \
"	int k;					 											 \n" \
"	int hsum;								 \n" \
"	int lsum;							 			 \n" \
"	loclsum[lid]=lowsum[id];						 			 \n" \
"	lochsum[lid] = highsum[id];							 			 \n" \
"	for (i=0; i<NUM_ITERATIONS; i++) {									 \n" \
"	 for(j = 0; j<lsize;j++) {  		   						 \n" \
"	  lsum=loclsum[j];								 \n" \
"	  hsum=lochsum[j];  								 \n" \
"	  for(k=lsum; k<hsum; k+=lsize){  								 \n" \
"	    if((k+lid)<hsum)  								 \n" \
"	      y[row[k+lid]]+=x[col[k+lid]]*val[k+lid];	 \n" \
"  	  }										 \n" \
"	 }	\n"
"	}	\n"
"}      \n" \
"}      \n" \
"\n";

////////////////////////////////////////////////////////////////////////////////
JNIEXPORT jfloatArray JNICALL Java_section2_sparsematmult_SparseMatmultOCL_sparseFromC
 (JNIEnv * env, jobject jobj, jint jnthreads, jfloatArray jval, jintArray jrow,jintArray jcol, jfloatArray jx, jfloatArray jy, jint jNUM_ITERATIONS, jintArray jlowsum, jintArray jhighsum, jint jnz_size, jint jn_size, jint jm_size)
{

    int err;                            // error code returned from api calls
     
    jfloatArray res;           			// results returned from device
   
    size_t global;                      // global domain size for our calculation
    size_t local;                       // local domain size for our calculation
    struct timeval tim;
    double starttime, endtime,temp;
    cl_device_id device_id;             // compute device id 
    cl_platform_id platform;		// compute platform id
    cl_context context;                 // compute context
    cl_command_queue commands;          // compute command queue
    cl_program program;                 // compute program
    cl_kernel kernel;                   // compute kernel
    
    cl_mem valb;
    cl_mem rowb;
    cl_mem colb;
    cl_mem xb;
    cl_mem yb;
    cl_mem lsumb;
    cl_mem hsumb;

	unsigned int nthreads = (unsigned int)jnthreads;
	unsigned int NUM_ITERATIONS = (unsigned int)jNUM_ITERATIONS;
	int nz_size = jnz_size;
	int n_size = jn_size;
	int m_size = jm_size;
	
   

	res = (*env)->NewFloatArray(env, m_size);
    	float* val = (*env)->GetFloatArrayElements(env, jval,0);
    	int* row = (*env)->GetIntArrayElements(env, jrow,0);
    	int* col = (*env)->GetIntArrayElements(env, jcol,0);
    	float* x = (*env)->GetFloatArrayElements(env, jx,0);
	float* y = (*env)->GetFloatArrayElements(env, jy,0);
	int* lowsum = (*env)->GetIntArrayElements(env, jlowsum,0);
	int* highsum = (*env)->GetIntArrayElements(env, jhighsum,0);
    



    // Connect to a compute device
   
    err = clGetPlatformIDs(1, &platform, NULL);
    if (err != CL_SUCCESS)
    {
        printf("Error: Failed to select paltform!\n");
        return res;
    }

	
    // Connect to a compute device
    //
    err = clGetPlatformIDs(3, &platform, NULL);
    if (err != CL_SUCCESS)
    {
        printf("Error: Failed to select paltform! %d\n",err);
        return res;
    }



    err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_GPU, 1, &device_id, NULL);
    if (err != CL_SUCCESS)
    {
        printf("Error: Failed to select device!\n");
        return res;
    }
    //
    context = clCreateContext(0, 1, &device_id, NULL, NULL, &err);
    if (!context)
    {
        printf("Error: Failed to create a compute context!\n");
        return res;
    }

    // Create a command queue
    //
    commands = clCreateCommandQueue(context, device_id, 0, &err);
    if (!commands)
    {
        printf("Error: Failed to create a command queue!\n");
        return res;
    }

    // Create the compute program from the source buffer
    //
    program = clCreateProgramWithSource(context, 1, (const char **) & KernelSource, NULL, &err);
    if (!program)
    {
        printf("Error: Failed to create compute program!\n");
        return res;
    }
 
    // Build the program executable
    //
    err = clBuildProgram(program, 0, NULL, NULL, NULL, NULL);
    if (err != CL_SUCCESS)
    {
        size_t len;
        char buffer[2048];
 
        printf("Error: Failed to build program executable!\n");
        clGetProgramBuildInfo(program, device_id, CL_PROGRAM_BUILD_LOG, sizeof(buffer), buffer, &len);
        printf("%s\n", buffer);
        exit(1);
    }

    // Create the compute kernel in the program we wish to run
    //
    kernel = clCreateKernel(program, "square", &err);
    if (!kernel || err != CL_SUCCESS)
    {
        printf("Error: Failed to create compute kernel!\n");
        exit(1);
    }


    // Create arrays in device memory for our calculations
    //
    yb = clCreateBuffer(context,  CL_MEM_READ_WRITE,  sizeof(float) * m_size, NULL, NULL);
    xb = clCreateBuffer(context, CL_MEM_READ_ONLY, sizeof(float) * n_size, NULL, NULL);
    valb = clCreateBuffer(context, CL_MEM_READ_ONLY, sizeof(float) * nz_size, NULL, NULL);
    rowb = clCreateBuffer(context, CL_MEM_READ_ONLY, sizeof(int) * nz_size, NULL, NULL); 
    colb = clCreateBuffer(context, CL_MEM_READ_ONLY, sizeof(int) * nz_size, NULL, NULL);
    lsumb = clCreateBuffer(context, CL_MEM_READ_ONLY, sizeof(int) * (nthreads+1), NULL, NULL); 
   	hsumb = clCreateBuffer(context, CL_MEM_READ_ONLY, sizeof(int) * (nthreads+1), NULL, NULL);

    if (!yb || !xb || !valb || !rowb || !colb || !lsumb || !hsumb)
    {
        printf("Error: Failed to allocate device memory!\n");
        exit(1);
    }    
    // Write our data sets into the arrays in device memory 
    //
    err = clEnqueueWriteBuffer(commands, yb, CL_TRUE, 0, sizeof(float) * m_size, y, 0, NULL, NULL);
    err |= clEnqueueWriteBuffer(commands, xb, CL_TRUE, 0, sizeof(float) * n_size, x, 0, NULL, NULL);
    err |= clEnqueueWriteBuffer(commands, valb, CL_TRUE, 0, sizeof(float) * nz_size, val, 0, NULL, NULL);
    err |= clEnqueueWriteBuffer(commands, lsumb, CL_TRUE, 0, sizeof(int) * (nthreads+1), lowsum, 0, NULL, NULL);
    err |= clEnqueueWriteBuffer(commands, hsumb, CL_TRUE, 0, sizeof(int) * (nthreads+1), highsum, 0, NULL, NULL);
    err |= clEnqueueWriteBuffer(commands, rowb, CL_TRUE, 0, sizeof(int) * nz_size, row, 0, NULL, NULL);   
    err |= clEnqueueWriteBuffer(commands, colb, CL_TRUE, 0, sizeof(int) * nz_size, col, 0, NULL, NULL);
    if (err != CL_SUCCESS)
    {
        printf("Error: Failed to write array to device memory! %d\n",err);
        exit(1);
    }

//get the number of compute units (SMs) of this device
    cl_uint units;
    clGetDeviceInfo( device_id, CL_DEVICE_MAX_COMPUTE_UNITS,  sizeof(cl_uint),  &units, NULL);

//get the multiple of local work group
    size_t mult;
    err = clGetKernelWorkGroupInfo(kernel, device_id, CL_KERNEL_PREFERRED_WORK_GROUP_SIZE_MULTIPLE, sizeof(mult), &mult, NULL);
    if (err != CL_SUCCESS)
    {
        printf("Error: Failed to retrieve kernel work group info! %d\n", err);
        exit(1);
    }
  
    global = nthreads;
    local = global / (units*32); // create at least 32 warps per multiprocessor, so the SM can execute some other warp while one of the warps waits (a group has at least one warp)
    if(local>mult) {
    local = local / mult; //now make local a multiple of mult (warp size) so that resources in the warp are not wasted
    local = local * mult;
    }

    
    if(global % local != 0)
    {
		global=global-(global%local);
		global+=local; //make global a multiple of local
    }

    // Set the arguments to our compute kernel
    //
    err = clSetKernelArg(kernel, 0, sizeof(cl_mem), &valb);
    err |= clSetKernelArg(kernel, 1, sizeof(cl_mem), &yb);
    err |= clSetKernelArg(kernel, 2, sizeof(cl_mem), &rowb);
    err |= clSetKernelArg(kernel, 3, sizeof(cl_mem), &colb);
    err |= clSetKernelArg(kernel, 4, sizeof(cl_mem), &xb);
    err |= clSetKernelArg(kernel, 5, sizeof(unsigned int), &NUM_ITERATIONS);
    err |= clSetKernelArg(kernel, 6, sizeof(unsigned int), &nthreads);
    err |= clSetKernelArg(kernel, 7, sizeof(cl_mem), &lsumb);
    err |= clSetKernelArg(kernel, 8, sizeof(cl_mem), &hsumb);
    err |= clSetKernelArg(kernel, 9, ((nthreads / (global/local))+1) *sizeof(int), NULL);
    err |= clSetKernelArg(kernel, 10, ((nthreads / (global/local))+1) *sizeof(int), NULL);
    if (err != CL_SUCCESS)
    {
        printf("Error: Failed to set kernel arguments! %d\n", err);
        exit(1);
    }


    // Execute the kernel
    //
    //measure the kernel time
    gettimeofday(&tim, NULL); 
    starttime =  tim.tv_sec*1000+(tim.tv_usec*0.001);
    err = clEnqueueNDRangeKernel(commands, kernel, 1, NULL, &global, &local, 0, NULL, NULL);
    if (err)
    {
        printf("Error: Failed to execute kernel!\n");
        return res;
    }
    

    // Wait for the commands to get serviced before reading back results
    //
    clFinish(commands);
    
    gettimeofday(&tim, NULL); 
    endtime  = tim.tv_sec*1000+(tim.tv_usec*0.001);
    temp = (endtime-starttime)*0.001;
    printf("%4.3f", temp);
    // Read back the res from the device
    //
      
    y = malloc(sizeof(float)*m_size);
    err = clEnqueueReadBuffer( commands, yb, CL_TRUE, 0, sizeof(float) * m_size, y, 0, NULL, NULL );  
    
    if (err != CL_SUCCESS)
    {
        printf("Error: Failed to read output array! %d\n", err);
        exit(1);
    }
    
    // Shutdown and cleanup
  
    (*env)->SetFloatArrayRegion(env,res, 0, m_size, y);
    free(y);
    clReleaseMemObject(valb);
    clReleaseMemObject(rowb);
    clReleaseMemObject(colb);
    clReleaseMemObject(xb);
    clReleaseMemObject(yb);
    clReleaseMemObject(lsumb);
    clReleaseMemObject(hsumb);
    clReleaseProgram(program);
    clReleaseKernel(kernel);
    clReleaseCommandQueue(commands);
    clReleaseContext(context);

    return res;
}
